package taskpool

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/zeromicro/go-zero/core/logx"
)

var (
	ErrPoolClosed    = errors.New("task pool is closed")
	ErrSubmitTimeout = errors.New("submit task timeout")
)

// TaskPool 任务池结构体
type TaskPool struct {
	workers       int           // 工作goroutine数量
	jobQueue      chan func()   // 任务队列
	poolCloseDone chan struct{} // 关闭信号
	closeOnce     sync.Once     // 确保关闭操作只执行一次
	wg            sync.WaitGroup
	errorHandler  func(error) // 错误处理函数
}

// NewTaskPool 创建任务池
// workers: 工作goroutine数量
// queueSize: 任务队列缓冲区大小
func NewTaskPool(workers, queueSize int) *TaskPool {
	if workers < 1 {
		workers = 1
	}

	p := &TaskPool{
		workers:       workers,
		jobQueue:      make(chan func(), queueSize),
		poolCloseDone: make(chan struct{}),
		errorHandler: func(err error) {
			logx.Errorf("task pool error", "error", err)
		},
	}

	p.start()
	return p
}

// 启动工作goroutine
func (p *TaskPool) start() {
	p.wg.Add(p.workers)
	for i := 0; i < p.workers; i++ {
		go p.worker()
	}
}

// 工作goroutine
func (p *TaskPool) worker() {
	defer p.wg.Done()

	for {
		select {
		case job, ok := <-p.jobQueue:
			if !ok {
				return
			}
			p.safeRun(job)
		case <-p.poolCloseDone:
			return
		}
	}
}

// 安全运行任务（处理Panic）
func (p *TaskPool) safeRun(job func()) {
	defer func() {
		if r := recover(); r != nil {
			if p.errorHandler != nil {
				err := fmt.Errorf("task panic: %v", r)
				p.errorHandler(err)
			}
		}
	}()

	job()
}

// Submit 提交任务
func (p *TaskPool) Submit(task func(), timeout time.Duration) error {
	if timeout < 0 {
		timeout = 0 // 负值视为无限等待
	}

	select {
	case <-p.poolCloseDone:
		return ErrPoolClosed
	default:
		return p.handleBufferedSubmit(task, timeout)
	}
}

// 处理带缓冲通道的提交
func (p *TaskPool) handleBufferedSubmit(task func(), timeout time.Duration) error {
	if timeout == 0 {
		select {
		case p.jobQueue <- task:
			return nil
		case <-p.poolCloseDone:
			return ErrPoolClosed
		default:
			return ErrSubmitTimeout
		}
	}

	select {
	case p.jobQueue <- task:
		return nil
	case <-p.poolCloseDone:
		return ErrPoolClosed
	case <-time.After(timeout):
		return ErrSubmitTimeout
	}
}

// SetErrorHandler 设置错误处理函数
func (p *TaskPool) SetErrorHandler(handler func(error)) {
	p.errorHandler = handler
}

// Close 优雅关闭任务池
func (p *TaskPool) Close() {
	p.closeOnce.Do(func() {
		close(p.poolCloseDone) // 先停止接收新任务
		close(p.jobQueue)      // 关闭任务队列，处理剩余任务
		p.wg.Wait()            // 等待所有工作goroutine退出
	})
}
